'''
https://leetcode.cn/problems/DFPeFJ/description/
'''
import heapq
from collections import defaultdict, deque
from curses.ascii import SO
from typing import List
import numpy as np


class Solution:
    def electricCarPlan(self, paths: List[List[int]], cnt: int, start: int, end: int, charge: List[int]) -> int:
        # paths->edges, cnt:(电车电池最大容量,能行驶的距离｜行驶cnt距离花费的时间), start,end:起始结束地点， charge: 每个城市充一个电所耗费的时间
        graph = defaultdict(list)
        for u, v, w in paths:
            graph[u].append((v, w))
            graph[v].append((u, w))
        n = len(charge)

        # 遍历图的每一个节点为：(哪个节点,当前的电量)
        distances = np.full((n, cnt + 1), np.inf)
        visited = np.full((n, cnt + 1), False)
        distances[start, 0] = 0

        heap = [(0, start, 0)]      # cur_cost, cur_city, cur_power
        while heap:
            cur_cost, u, cur_power = heapq.heappop(heap)
            if visited[u, cur_power]: continue
            if u == end: return cur_cost
            visited[u, cur_power] = True
            # 当前位置充一格电，和 去下一个城市
            if cur_power < cnt:
                next_cost = cur_cost + charge[u]
                v = u
                next_power = cur_power + 1
                if next_cost < distances[v, next_power]:
                    distances[v, next_power] = next_cost
                    heapq.heappush(heap, (next_cost, v, next_power))
            for v, w in graph[u]:
                if cur_power < w: continue  # 当前电量不支持去这个城市
                next_cost = cur_cost + w
                next_power = cur_power - w
                if next_cost < distances[v, next_power]:
                    distances[v, next_power] = next_cost
                    heapq.heappush(heap, (next_cost, v, next_power))
        return -1

paths = [[1,3,3],[3,2,1],[2,1,3],[0,1,4],[3,0,5]]
cnt = 6
start = 1
end = 0
charge = [2,10,4,1]
print(Solution().electricCarPlan(paths, cnt, start, end , charge))